import os
import paramiko
from pathlib import Path
import traceback

def find_working_directory(sftp, root_dir, prefix_to_match):
    """使用SFTP查找包含指定前缀文件的工作目录"""
    try:
        # 列出根目录下所有子目录
        for dir_name in sftp.listdir(root_dir):
            print(f"check dir {dir_name}")
            if not (dir_name.count('_') >= 2):  # 匹配*_*_*模式
                continue
            
            dir_path = f"{root_dir}/{dir_name}"
            a_cat_path = f"{dir_path}/A_cat"
            try:
                # 检查A_cat目录是否存在
                sftp.listdir(a_cat_path)
                
                # 检查A_cat目录下是否存在匹配文件
                for file_name in sftp.listdir(a_cat_path):
                    if file_name.startswith(prefix_to_match):
                        return Path(dir_path)
                        
            except IOError:
                continue  # A_cat目录不存在
                
    except Exception as e:
        print(f"查找工作目录时出错: {e}")
    return None
def copy_files(ssh_client, sftp, working_dir, target_prefix):
    """使用SFTP执行所有文件复制操作"""
    saveDir = f"data/{target_prefix}"
    os.makedirs(saveDir, exist_ok=True)  # 确保本地目标目录存在
    
    try:
        
        # 主文件复制
        files_to_copy = [
            # ('F_diff', f'{target_prefix}_cmb005.fit', f'{target_prefix}_cmb005_resi.fit'),
            # ('E_diff_align', f'{target_prefix}_cmb005.fit', f'{target_prefix}_cmb005_align.fit'),
            # ('D_combine', f'{target_prefix}_cmb005.fit', f'{target_prefix}_cmb005_obj.fit'),
            ('G_diff_cat', f'{target_prefix}_cmb005_aligned_wcs.cat', None),
            ('G_diff_cat', f'{target_prefix}_cmb005_aligned_wcs_unmch_template.cat', None),
            ('G_diff_cat', f'{target_prefix}_cmb005_aligned_wcs_unmch_tmpl_ucac4.cat', None),
            ('G_diff_cat', f'{target_prefix}_cmb005_aligned_wcs_unmch_tmpl_ucac4_hpunmch.cat', None),
            ('G_diff_cat', f'{target_prefix}_cmb005_aligned_wcs_unmch_tmpl_ucac4_hpunmch_mpmch.cat', None),
            ('G_diff_cat', f'{target_prefix}_cmb005_aligned_wcs_unmch_tmpl_ucac4_hpunmch_mpunmch.cat', None),
            ('G_diff_cat', f'{target_prefix}_cmb005_aligned_wcs_unmch_tmpl_ucac4_hpcorr.cat', None),
            ('G_diff_cat', f'{target_prefix}_cmb005_pred.cat', None),
            ('G_diff_cat', f'{target_prefix}_cmb005_blindmch_filtered_correct.cat', None),
            ('G_diff_cat', f'{target_prefix}_cmb005_blindmch_filtered_correct.reg', None),
        ]
        
        for src_dir, src_file, dst_file in files_to_copy:
            src_path = f'{working_dir}/{src_dir}/{src_file}'
            dst_name = dst_file if dst_file else src_file
            dst_path = f'{saveDir}/{dst_name}'
            print(f"复制文件: {src_path} 到 {dst_path}")
            try:
                # 直接通过SFTP复制文件
                sftp.get(src_path, dst_path)  # 下载到本地临时文件
            except Exception as e:
                print(f"复制{src_path}出错,错误信息: {e}")
                # traceback.print_exc()
        
        # WCS模板处理
        stdin, stdout, stderr = ssh_client.exec_command(f'ls {working_dir}/H_wcs_template/')
        wcs_subdirs = stdout.read().decode().split()
        
        if len(wcs_subdirs) == 1:
            wcs_subdir = wcs_subdirs[0]
            wcs_files = [
                ('hptemp.txt', 'hptemp.txt'),
                # ('template.fits', 'template.fits'),
                ('template.cat', 'template.cat'),
                ('template_4detection.cat', 'template_4detection.cat'),
                ('bright_star_align_filtered.cat', 'bright_star_align_filtered.cat'),
                ('bright_star_align_filtered_correct.cat', 'bright_star_align_filtered_correct.cat'),
                ('bright_star_align_filtered_correct.reg', 'bright_star_align_filtered_correct.reg'),
                ('bright_star_deep_align_filtered.cat', 'bright_star_deep_align_filtered.cat'),
                ('bright_star_deep_align_filtered_correct.cat', 'bright_star_deep_align_filtered_correct.cat'),
                ('bright_star_deep_align_filtered_correct.reg', 'bright_star_deep_align_filtered_correct.reg'),
                ('mp_pred_template.dat', 'mp_pred_template.dat')
            ]
            
            for src, dst in wcs_files:
                src_path = f'{working_dir}/H_wcs_template/{wcs_subdir}/{src}'
                dst_path = f'{saveDir}/{dst}'
                
                try:
                    sftp.get(src_path, dst_path)  # 下载到本地临时文件
                except Exception as e:
                    print("*"*20)
                    print(f"复制{src_path}出错,错误信息: {e}")
                    
    except Exception as e:
        print(f"复制文件过程中发生错误: {e}")
    finally:
        sftp.close()

    remoteIP = "10.36.1.155"
    print(f"开始将数据复制到远程服务器gwac@{remoteIP}")
    try:
        remote_ssh = paramiko.SSHClient()
        remote_ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        # remote_ssh.connect('10.34.1.111', username='gwac', password='gwac1234')
        remote_ssh.connect(remoteIP, port=9901, username='gwac', password='gwac1234', timeout=10)
        remote_sftp = remote_ssh.open_sftp()

        # 获取用户主目录路径
        remote_home = remote_sftp.normalize('.')  # 通常为 /home/gwac
        remote_target_dir = f"{remote_home}/{target_prefix}"

        try:
            remote_sftp.mkdir(remote_target_dir)
        except IOError:
            pass  # 目录可能已存在

        for file_name in os.listdir(saveDir):
            local_path = os.path.join(saveDir, file_name)
            remote_path = f"{remote_target_dir}/{file_name}"
            print(f"from {local_path} to {remote_path}")
            remote_sftp.put(local_path, remote_path)
            print(f"已上传 {file_name} 到远程服务器")

        print("所有文件已成功复制到远程服务器")
    except Exception as e:
        print(f"复制到远程服务器时出错: {e}")
        traceback.print_exc()
    finally:
        try:
            remote_sftp.close()
            remote_ssh.close()
        except:
            pass
def main():
    # 配置参数
    remote_user = 'gwac'
    remote_ip = '172.28.2.41'
    remote_pass = 'gwac1234'  # 请替换为实际密码
    
    root_dir = '/data/work/gwac_diff_data_new/20251002'
    #G061_Mon_objt_250120T14165670
    prefix_to_match = 'G041_Mon_objt_250901T19453459'
        
    # 添加日志记录
    print("="*50)
    print("开始执行文件复制任务")
    print(f"目标服务器: {remote_user}@{remote_ip}")
    print(f"根目录: {root_dir}")
    print(f"目标前缀: {prefix_to_match}")
    print("="*50)
    
    # 初始化SSH客户端
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    
    try:
        ssh.connect(remote_ip, username=remote_user, password=remote_pass)
        sftp = ssh.open_sftp()
        
        # 查找工作目录
        working_dir = find_working_directory(sftp, root_dir, prefix_to_match)
        if not working_dir:
            print(f"未找到包含 {prefix_to_match} 的工作目录")
            return
        
        print(f"找到工作目录: {working_dir}")
        
        # 执行文件复制
        copy_files(ssh, sftp, str(working_dir), prefix_to_match)
        
        print("="*50)
        print("文件复制任务完成")
        print(f"所有文件已复制到远程服务器的 ~/{prefix_to_match} 目录")
        print("="*50)
            
    except Exception as e:
        print(f"SSH连接或操作过程中发生错误: {e}")
    finally:
        try:
            sftp.close()
        except:
            pass
        ssh.close()

if __name__ == '__main__':
    main()
